import numpy as np
from sklearn.metrics import r2_score

class Lasso():
    def __init__(self):
        pass

    # 读取数据
    def prepare_data(self):
        data = np.genfromtxt('./example.dat', delimiter=',')
        x = data[:, 0:100]
        y = data[:, 100].reshape(-1, 1)
        # 在x前面加一列
        x = np.column_stack((np.ones((x.shape[0], 1)), x))
        # 数据划分
        x_train, y_train = x[:70], y[:70]
        x_test, y_test = x[70:], y[70:]
        return x_train,y_train,x_test,y_test

    # 初始化w， b参数
    def initialize(self, dims):
        w = np.zeros((dims, 1))
        b = 0
        return w, b
    def sign(self,x):
        if x > 0:
            return 1
        elif x < 0:
            return -1
        else:
            return 0

    # 损失函数
    def l1_loss(self,X,y,w,b,alpha):
        num_train = X.shape[0]
        num_feature = X.shape[1]

        y_hat = np.dot(X,w) + b
        loss = np.sum((y_hat-y)**2) / num_train + np.sum(alpha*abs(w))
        dw = np.dot(X.T, (y_hat-y)) / num_train + alpha*np.vectorize(self.sign)(w)
        db = np.sum((y_hat-y)) / num_train
        return y_hat, loss, dw, db
    # 训练函数
    def lasso_train(self,X, y, w, b, learning_rate, epochs):
        loss_list = []
        w, b = self.initialize(X.shape[1])
        for i in range(1, epochs):
            y_hat, loss, dw, db = self.l1_loss(X, y, w, b, 0.1)
            w += -learning_rate * dw
            b += -learning_rate *db
            loss_list.append(loss)

            if i % 50 == 0:
                print('epoch %d loss %f' %(i, loss))

            params = {
                'w':w,
                'b':b
            }
            grads = {
                'dw':dw,
                'db':db
            }
        return loss, loss_list, params, grads
    def predict(self, x, params):
        w = params['w']
        b = params['b']
        y_pred = np.dot(x, w) + b
        return y_pred

    # 主程序序
if __name__ == '__main__':
    lasso = Lasso()
    X_train, y_train, X_test, y_test = lasso.prepare_data()
    w, b = lasso.initialize(X_train.shape[0])
    loss, loss_list, params, grads = lasso.lasso_train(X_train, y_train, w, b, 0.01, 1000)
    print(params)
    y_pred = lasso.predict(X_test, params)
    print(r2_score(y_test, y_pred))



